package com.zenithmind.user.security;

import org.springframework.core.convert.converter.Converter;
import org.springframework.security.authentication.AbstractAuthenticationToken;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.authority.SimpleGrantedAuthority;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.server.resource.authentication.JwtAuthenticationToken;
import org.springframework.security.oauth2.server.resource.authentication.JwtGrantedAuthoritiesConverter;
import org.springframework.stereotype.Component;

import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * Converter to extract roles from JWT token and convert to Spring Security authorities
 */
@Component
public class JwtAuthenticationConverter implements Converter<Jwt, AbstractAuthenticationToken> {

    private final JwtGrantedAuthoritiesConverter defaultConverter = new JwtGrantedAuthoritiesConverter();

    @Override
    public AbstractAuthenticationToken convert(Jwt jwt) {
        // Default authorities from scope
        Collection<GrantedAuthority> defaultAuthorities = defaultConverter.convert(jwt);
        
        // Extract roles from the roles claim
        Collection<GrantedAuthority> customAuthorities = extractRoles(jwt);
        
        // Combine authorities
        Collection<GrantedAuthority> allAuthorities = Stream.concat(
                defaultAuthorities != null ? defaultAuthorities.stream() : Stream.empty(),
                customAuthorities.stream())
            .collect(Collectors.toList());
            
        return new JwtAuthenticationToken(jwt, allAuthorities, getPrincipalClaimName(jwt));
    }
    
    private String getPrincipalClaimName(Jwt jwt) {
        // Use "preferred_username" if available, fallback to "sub"
        String username = jwt.getClaim("preferred_username");
        if (username == null) {
            username = jwt.getSubject();
        }
        return username;
    }
    
    private Collection<GrantedAuthority> extractRoles(Jwt jwt) {
        Object roles = jwt.getClaim("roles");
        
        if (roles == null) {
            return Collections.emptyList();
        }
        
        Collection<String> rolesList;
        if (roles instanceof String) {
            // Single role as string
            rolesList = Arrays.asList(((String) roles).split(","));
        } else if (roles instanceof Collection) {
            // Collection of roles
            rolesList = (Collection<String>) roles;
        } else {
            return Collections.emptyList();
        }
        
        // Convert roles to authorities with ROLE_ prefix
        return rolesList.stream()
            .map(role -> "ROLE_" + role)
            .map(SimpleGrantedAuthority::new)
            .collect(Collectors.toList());
    }
} 